'''
Functions for generating random mathematical objects.
'''
import random
import sympy as sp
import functools
from .numeric import INTEGER_GEN, RATIONAL_GEN
from ..complexity import complexity
One = sp.Integer(1)

#===============================================================================
# Public API
#===============================================================================
__all__ = []
def public(func):
    'Decorator that makes the function public'
    __all__.append(func.__name__)
    return func

def multiple_generator(gen, *, args=(), kwds={}, size=None, unique=True):
    idx = 0
    objects = set()
    while size is None or idx < size:
        obj = gen(*args, **kwds)
        if unique:
            if obj not in objects:
                objects.add(obj)
            else:
                continue
        idx += 1
        yield obj

#===============================================================================
# Numeric generators
#===============================================================================
@public
def r_sign(p=0.5):
    '''
    Return either 1 or -1. ``p`` is the probability of returning a positive 
    value.
    '''
    if p >= random.random():
        return One
    else:
        return -One

def numeric_r_api(func):
    '''Decorator for numeric generators that implements the ``force_pos`` and 
    ``zero`` arguments. It also calls the decorated function with ``complexity``
    set to tuple argument, so implementations can expect this uniform calling
    convention.    
    '''
    @functools.wraps(func)
    def decorated(complexity=3, force_pos=True, zero=False):
        # 5% chance of creating zeros if zero is True
        if zero and 0.05 > random.random():
            return One * 0

        # Complexity can be a maximum value or a bound (when it is a 2-tuple)
        try:
            a, b = complexity
        except TypeError:
            a, b = 1, complexity

        # Creates random number
        out = func((a, b), force_pos, zero)

        # Change sign, if necessary
        if not force_pos:
            out *= r_sign()
        return One * out

    return decorated

def numeric_gen_api(func):
    @functools.wraps(func)
    def decorated(complexity=5, force_pos=True, zero=False, init=None, size=None, unique=True):
        # Check if init values are unique and yield them
        idx = 0
        init = list(init or [])
        if unique and len(set(init)) != len(init):
            raise ValueError('Initial elements are not unique')
        if len(init) > size:
            for obj in init[:size]:
                yield obj
            return
        else:
            for obj in init:
                yield obj
            idx += len(init)

        # Generate non-unique values if unique=False
        if not unique:
            for _ in range(size - idx):
                yield func(complexity, force_pos, zero)
            return

        # Create a set of values to check uniqueness
        values = set(init)
        for _ in range(size - idx):
            # 100 attempts to create a unique value
            for _ in range(100):
                new = func(complexity, force_pos, zero)
                if new in values:
                    values.add(new)
                    yield new
                    break
            else:
                raise ValueError('too many failed attempts to generate a unique number')
    return decorated

# Integers ---------------------------------------------------------------------
@public
@numeric_r_api
def r_integer(complexity=3, force_pos=True, zero=False):
    '''
    Returns a random integer of given complexity.
    
    Arguments
    ---------
    
    complexity: int, 2-tuple
        Specify either the maximum complexity or the complexity bound of the 
        random numbers generated by this function.
    force_pos: bool
        If True, returns only positive numbers.
    zero : bool
        If True, it can return 0.
    '''
    a, b = complexity
    return INTEGER_GEN.new(a, b)

@public
def r_int(complexity=3, force_pos=True, zero=False):
    '''Alias to r_integer()'''

    return r_integer(complexity, force_pos, zero)

# Rational numbers -------------------------------------------------------------
@public
@numeric_r_api
def r_rational(complexity=3, force_pos=True, zero=False):
    '''
    Returns a random rational number of given complexity.
    
    Arguments: see `r_integer()`
    '''
    a, b = complexity
    return RATIONAL_GEN.new(a, b)

@public
def r_frac(complexity=3, force_pos=True, zero=False):
    '''Alias to r_rational()'''

    return r_rational(complexity, force_pos, zero)


@public
def r_rat(complexity=3, force_pos=True, zero=False):
    '''Alias to r_rational()'''

    return r_rational(complexity, force_pos, zero)

@public
@numeric_gen_api
def gen_rational(complexity=5, force_pos=True, zero=False, init=None, size=None, unique=True):
    '''
    Iterator over a (possibly infinite) sequence of random fractions. 
    
    Examples
    --------
    
    >>> lst = list(gen_rational(init=[2, One/2], size=4, complexity=5))
    >>> lst[0] == 2 and lst[1] == One/2
    True
    >>> Integer(lst[3]).complexity() <= 5
    True
    '''
    return r_rational(complexity, force_pos, zero)

# Real numbers -----------------------------------------------------------------
# TODO

# Complex numbers --------------------------------------------------------------
# TODO

#===============================================================================
# Query functions
#===============================================================================
@public
def issimple(obj, max_complexity=5):
    '''
    Return True if object's complexity is less then or equal to ``max_complexity``
    
    Examples
    --------
    
    >>> issimple(6)
    True
    >>> issimple(6, 2)
    False
    '''
    return round(complexity(obj)) <= max_complexity

@public
def assure_simple(obj, max_complexity=5):
    '''
    Raises an AssertionError if object's complexity is greater than 
    ``max_complexity``.
    '''
    compl = round(complexity(obj))
    if not(compl <= max_complexity):
        msg = 'complexity of %s greater than the maximum complexity %s'
        msg = msg % (compl, max_complexity)
        raise AssertionError(msg)


#
#def rnumbers(max_compl=5, force_pos=False, init=None, max_size=None, unique=True):
#    '''
#    Iterator over a (possibly infinite) sequence of random integers. 
#    
#    Examples
#    --------
#    
#    >>> lst = list(rnumbers(init=[1, 2], max_size=4, max_compl=5))
#    >>> lst[0] == 1 and lst[1] == 2
#    True
#    >>> Integer(lst[3]).complexity() <= 5
#    True
#    '''
#
#    init = ([] if init is None else map(sp.sympify, init))
#    max_size = (2 ** 100 if max_size is None else max_size)
#    values = set(init)
#    values_lst = list(init) if init else [1]
#    if unique and (len(values) != len(init)):
#        raise ValueError('Initial elements are not unique')
#
#    # Yield values of the initial list
#    iter = 0
#    for x in init:
#        iter += 1
#        if iter > max_size:
#            return
#        yield x
#
#    # Create new numbers randomly
#    while True:
#        new = random.choice(values_lst)
#        new *= rnumber(max_compl, force_pos)
#        if new in values:
#            continue
#        else:
#            compl = Integer(new).complexity()
#            if compl > max_compl:
#                continue
#            else:
#                iter += 1
#                if iter > max_size:
#                    return
#                values.add(new)
#                values_lst.append(new)
#                yield new
#

if __name__ == '__main__':
    import doctest
    doctest.testmod()

lst = list(gen_rational(init=[2, One / 2], size=4, complexity=5))
